from luaparser import astnodes
from luaparser import ast


def main():
    lua_filepath = "./example.lua"
    with open(lua_filepath) as f:
        chunk = ast.parse(f.read())
    ast_path = "./tree.ast"
    with open(ast_path, "w") as f:
        f.write(ast.to_pretty_str(chunk))

class ReturnValue(Exception):
    pass

class LuaParser:
    def __init__(self, globals_={}):
        self.variable_stack = []
        self.locals = {}
        self.locals.update(globals)
        self.variable_stack.append(self.locals)
        self.return_value = None

    def get_value(self, name: str):
        for i in range(self.variable_stack.len()-1, 0, -1):
            value = self.variable_stack[i].get(name)
            if value:
                return value
        else:
            raise Exception(f"Not Found identifier: {name}")
        
    def get_var_stack_level(self, name: str):
        for i in range(self.variable_stack.len()-1, 0, -1):
            value = self.variable_stack[i].get(name)
            if value:
                return i
        return -1
    
    def get_logical_value(self, node: astnodes.Expression):
        return not (self.visit(node) in (False, None))

    def visit(self, node):
        assert isinstance(node, astnodes.Node), f"{node} not a Node"
        method = getattr(self, f"visit_{node.name}")
        return method(self, node)

    def visit_Chunk(self, node: astnodes.Chunk):
        block = node.body  # always a Block
        return self.visit(block)
    
    def visit_Block(self, node: astnodes.Block):
        self.locals = {}
        self.variable_stack.append(self.locals)
        try:
            for stmt in node.body:
                self.visit(stmt)
        finally:
            self.variable_stack.pop()
    
    def visit_LocalAssign(self, node: astnodes.LocalAssign):
        # 未考虑左值和右值中"..."的情况
        targets = node.targets  # always Name?
        values = node.values
        for idx, value_node in enumerate(values):
            target = targets[idx]
            name = target.id
            self.locals[name] = self.visit(value_node)
    
    def visit_Assign(self, node: astnodes.Assign):
        targets = node.targets
        values = node.values
        for idx, value_node in enumerate(values):
            target = targets[idx]
            value = self.visit(value_node)
            if isinstance(target, astnodes.Index):
                # lua table元素的赋值
                self.visit(target.value)[target.idx.id] = value
            elif isinstance(target, astnodes.Name):
                name = target.id
                level = self.get_var_stack_level(name)
                if level != -1:
                    # 现有变量的值修改
                    self.variable_stack[level][name] = value
                else:
                    # 新声明全局变量
                    self.variable_stack[0][name] = value
    
    def visit_Table(self, node: astnodes.Table):
        # list or dict 先尝试总是dict
        table = dict()
        for field in node.fields:
            if isinstance(field.key, astnodes.Number):
                k = self.visit(field.key)
            elif isinstance(field.key, astnodes.Name):
                k = field.key.id
            else:
                raise Exception("Unexpected field key type {}", type(field.key))
            v = self.visit(field.value)
            table[k] = v
        return table
    
    def visit_Number(self, node: astnodes.Number):
        return node.n
    
    def visit_String(self, node: astnodes.String):
        # delimiter是否影响解析?
        return node.s
    
    def visit_AddOp(self, node: astnodes.AddOp):
        return self.visit(node.left) + self.visit(node.right)
    
    def visit_SubOp(self, node: astnodes.SubOp):
        return self.visit(node.left) - self.visit(node.right)
    
    def visit_MultOp(self, node: astnodes.MultOp):
        return self.visit(node.left) * self.visit(node.right)
    
    def visit_FloatDivOp(self, node: astnodes.FloatDivOp):
        return self.visit(node.left) / self.visit(node.right)
    
    def visit_FloorDivOp(self, node: astnodes.FloorDivOp):
        return self.visit(node.left) // self.visit(node.right)
    
    def visit_ModOp(self, node: astnodes.ModOp):
        return self.visit(node.left) % self.visit(node.right)
    
    def visit_ExpoOp(self, node: astnodes.ExpoOp):
        return self.visit(node.left) ** self.visit(node.right)
    
    def visit_RLtOp(self, node: astnodes.LessThanOp):
        return self.visit(node.left) < self.visit(node.right)
    
    def visit_RGtOp(self, node: astnodes.GreaterThanOp):
        return self.visit(node.left) > self.visit(node.right)
    
    def visit_RLtEqOp(self, node: astnodes.LessOrEqThanOp):
        return self.visit(node.left) <= self.visit(node.right)
    
    def visit_RGtEqOp(self, node: astnodes.GreaterOrEqThanOp):
        return self.visit(node.left) >= self.visit(node.right)

    def visit_REqOp(self, node: astnodes.EqToOp):
        # Lua 有没有NaN的情况?
        return self.visit(node.left) == self.visit(node.right)

    def visit_RNotEqOp(self, node: astnodes.NotEqToOp):
        return self.visit(node.left) != self.visit(node.right)
    
    def visit_LAndOp(self, node: astnodes.AndLoOp):
        return self.get_logical_value(node.left) and self.get_logical_value(node.right)
    
    def visit_LOrOp(self, node: astnodes.OrLoOp):
        return self.get_logical_value(node.left) or self.get_logical_value(node.right)
    
    def visit_Concat(self, node: astnodes.Concat):
        return self.visit(node.left) + self.visit(node.right)
    
    def visit_UminusOp(self, node: astnodes.UMinusOp):
        return - self.visit(node.operand)
    
    def visit_ULNotOp(self, node: astnodes.ULNotOp):
        return not self.get_logical_value(node.operand)
        
    def visit_Call(self, node: astnodes.Call):
        func_node = self.visit(node.func)
        parameters = func_node.args
        # 未考虑参数收集和参数展开
        self.locals = {}
        self.variable_stack.append(self.locals)
        for idx, arg in enumerate(node.args):
            param = parameters[idx].id
            self.locals[param] = self.visit(arg)
        try:
            self.visit(func_node.body)
        except ReturnValue:
            pass
        finally:
            self.variable_stack.pop()
            result = self.return_value
            self.return_value = None
            return result
           
    def visit_LocalFunction(self, node: astnodes.LocalFunction):
        fn_name = node.name.id
        # 保留完整函数定义 到调用时才能传递参数展开计算函数内语句
        # 尚未考虑闭包 如若考虑还应保留当前的变量环境到该函数对象下
        self.locals[fn_name] = node

    def visit_Function(self, node: astnodes.Function):
        fn_name = node.name.id
        # 保留完整函数定义 到调用时才能传递参数展开计算函数内语句
        # 尚未考虑闭包 如若考虑还应保留当前的变量环境到该函数对象下
        self.variable_stack[0][fn_name] = node

    def visit_Index(self, node: astnodes.Index):
        value = self.visit(node.value)
        if node.notation == astnodes.IndexNotation.DOT:
            # dot 方式访问必定是Name 直接访问id 用visit访问反而尝试将字段名作为变量名解析
            return value[node.idx.id]
        elif node.notation == astnodes.IndexNotation.SQUARE:
            pass
        else:
            raise Exception("not a expected Notation Type: {}", type(node.notation))

    def visit_Name(self, node: astnodes.Name):
        id_ = node.id
        return self.get_value(id_)
    
    def visit_Invoke(self, node: astnodes.Invoke):
        source = self.visit(node.source)
        func = getattr(source, node.func.id)
        args = tuple(map(self.visit, node.args))
        return func(source, *args)

    def visit_If(self, node: astnodes.If):
        # else if 语句尚未处理
        if self.get_logical_value(node.test):
            self.visit(node.body)
        else:
            if node.orelse:
                self.visit(node.orelse)

    def visit_Return(self, node: astnodes.Return):
        # 未处理多返回值
        self.return_value = node.values[0] if node.values else None
        raise ReturnValue

    def visit_Nil(self, node: astnodes.Nil):
        return None
    
    def visit_True(self, node: astnodes.TrueExpr):
        return True
    
    def visit_False(self, node: astnodes.FalseExpr):
        return False

LUA_BUILTINS = {
    # "true": True,
    # "false": False,
}


if __name__ == "__main__":
    main()